#include <gtest/gtest.h>
#include <vector>
#include <opencv2/opencv.hpp>

#include "trt_define.h"
#include "utils/logger.h"
#include "utils/file_utils.h"
#include "utils/nv_utils.h"
#include "trt_engine/trt_tensor.h"
#include "trt_builder/trt_builder.h"
#include "app_yolo/yolo.h"
#include "app_yolo/object_detector.hpp"
#include "utils/elapsed_timer.h"


_TRT_INFER_BEGIN

class YoloTest : public ::testing::Test
{
public:
    const std::string test_dir = "D:\\learnSpace\\trt_infer\\test_data";
    const std::string engine_file = FileUtils::concat_path(test_dir, "yolov5s_fp16.engine");
    const std::string onnx_file = FileUtils::concat_path(test_dir, "yolov5s.onnx");
    // const std::string engine_file = FileUtils::concat_path(test_dir, "yolox_tiny_8x8_300e_coco.engine");
    // const std::string onnx_file = FileUtils::concat_path(test_dir, "yolox_tiny_8x8_300e_coco.onnx");
    const std::string image_file = FileUtils::concat_path(test_dir, "dog.jpg");
protected:
    YoloTest() {}
    ~YoloTest() override {
    }
    void SetUp() override {
    }
    void TearDown() override {
    }
};

TEST_F(YoloTest, build_trt_model)
{
    LOGER_INST.init();
    EXPECT_TRUE(FileUtils::is_file_exist(image_file));
    EXPECT_TRUE(FileUtils::is_file_exist(onnx_file));
    if (!FileUtils::is_file_exist(engine_file)) 
    {
        LOGI << "Trt engine file not exist, build it from onnx file";
        TrtBuilder::compile(TrtBuilder::ModelType::FP16, 10, onnx_file, engine_file, 1ul << 30);
    }
}

TEST_F(YoloTest, Box)
{
    Yolo::Box box(0, 0, 100, 100, 0.5, 1);
    EXPECT_EQ(box.area(), 10000);

    Yolo::Box box1(0, 0, 50, 50, 0.5, 1);
    Yolo::Box box2(25, 25, 75, 75, 0.5, 1);
    auto box3 = box1 | box2;
    auto box4 = box1 & box2;
    EXPECT_EQ(box3, Yolo::Box(0, 0, 75, 75, 0.5, 1));
    EXPECT_EQ(box4, Yolo::Box(25, 25, 50, 50, 0.5, 1));
    EXPECT_FLOAT_EQ(box2.iou(box1), 0.1428571);

    Yolo::BoxArray boxes = {{0, 0, 100, 100, 0.2, 1}, {25, 25, 75, 75, 0.5, 1}, {50, 50, 100, 100, 0.8, 1}};
    auto res = Yolo::nms_cpu(boxes, 0.2);
    EXPECT_EQ(res.size(), 2);
    EXPECT_EQ(res[0], Yolo::Box(50, 50, 100, 100, 0.8, 1));
    EXPECT_EQ(res[1], Yolo::Box(25, 25, 75, 75, 0.5, 1));
}

TEST_F(YoloTest, inference)
{
    LOGI << "Inference";
    EXPECT_TRUE(FileUtils::is_file_exist(engine_file));
    auto yolo_infer = Yolo::create_infer(
        engine_file, Yolo::Type::V5, 0, 0.25, 0.45
    );
    EXPECT_TRUE(yolo_infer);
    const std::string bus = "bus.jpg";
    const std::string bus_image = FileUtils::concat_path(test_dir, bus);
    EXPECT_TRUE(FileUtils::is_file_exist(bus_image));
    cv::Mat image = cv::imread(bus_image);
    EXPECT_FALSE(image.empty());
    auto boxes = yolo_infer->commit(image).get();
    EXPECT_GT(boxes.size(), 0);
    for (auto box : boxes)
    {
        LOGI << box;
        cv::rectangle(image, cv::Point(box.left, box.top), cv::Point(box.right, box.bottom), cv::Scalar(0, 255, 0), 2);
        cv::imwrite(FileUtils::concat_path(test_dir, FileUtils::concat_path(test_dir, "dog.jpg")), image);
    }

}


_TRT_INFER_END